import argparse
import os
from functools import reduce

import cv2

from utils.list import create_file_list
from utils.revolver import Revolver


def parse_args():
    parser = argparse.ArgumentParser(
        description='predictions parameters')
    parser.add_argument(
        '--root',
        dest='root',
        help='image to use',
        default='/home/admins/data/beer_data',
        type=str)
    parser.add_argument(
        '--image-path',
        dest='image_path',
        help='path of image to use in root, split by .',
        default='',
        type=str)
    parser.add_argument(
        '--output-root',
        dest='output',
        help='output root',
        default='',
        type=str)
    return parser.parse_args()


def get_revolved_image(params):
    image_lists = params['image_list']
    output_root = params['output_root']
    for idx, image in enumerate(image_lists):
        print('revolving amd cropping {} of {} images'.format(idx, len(image_lists)))
        revolver = Revolver(image)
        revolver.update()
        images = revolver.output
        out_root = os.path.join(output_root, '{:04}'.format(idx // 1000),
                                '{:08}'.format(idx))
        if not os.path.exists(out_root):
            os.makedirs(out_root)
        for _, angle in images.items():
            for _, seat in angle.items():
                for key, value in seat.items():
                    cv2.imwrite(os.path.join(out_root, '{}.jpg'.format(key)), value['image'])


def process():
    params = {}
    image_root = reduce(lambda x, y: os.path.join(x, y),
                        args.image_path.split('.'), args.root)

    def _filter(x, y, _):
        if x.split('.')[-1] in ['jpg', 'png', 'JPG', 'PNG']:
            y[0].append(x)
        return y

    image_lists, _ = create_file_list(image_root, filtering=_filter)
    output_root = args.root if args.output == '' else args.output
    if not os.path.exists(output_root):
        os.makedirs(output_root)
    params['image_list'] = image_lists
    params['output_root'] = output_root
    get_revolved_image(params)


if __name__ == '__main__':
    args = parse_args()
    process()
